package com.lxg.ai.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

@RestController
@RequestMapping("/embedding")
public class EmbeddingController {

    @Autowired
    private QdrantClient qdrantClient;

    @Autowired
    private EmbeddingModel embeddingModel;

    /**
     * 创建向量集合
     */
    @GetMapping("createCollection")
    public void createCollection(@RequestParam String collectionName){
        Collections.VectorParams vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)  //
                .build();
        qdrantClient.createCollectionAsync(collectionName, vectorParams);
    }


    /**
     * 保存向量文本
     * @param collectionName
     */
    @GetMapping("createTextSegment")
    public void createTextSegment(@RequestParam String collectionName){
        QdrantEmbeddingStore embeddingStore = QdrantEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6334)
                .collectionName(collectionName)
                .build();
        TextSegment segment = TextSegment.from("浏览器报错 404，请检测您输入的路径是否正确");
        segment.metadata().put("author", "冷冷");
        Embedding embedding = embeddingModel.embed(segment).content();
        embeddingStore.add(embedding, segment);
        embeddingStore.close();
    }


    /**
     * 向量查询
     */
    @GetMapping("getTextSegment")
    public void getTextSegment(@RequestParam String collectionName){
        QdrantEmbeddingStore embeddingStore = QdrantEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6334)
                .collectionName(collectionName)
                .build();
        Embedding queryEmbedding = embeddingModel.embed("404 是哪里的问题？").content();


        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
//        System.out.println(searchResult.matches().get(0).score());  // 输出匹配分数
//        System.out.println(searchResult.matches().get(0).embedded().text());  // 输出匹配文本
        System.out.println(searchResult.matches().get(0).embedded().text());
    }

    /**
     * 常用过滤器：
     * Filter名称	功能	使用示例
     * And	同时满足多个条件	Filter.and(condition1, condition2)
     * Or	满足其中任意一个条件	Filter.or(condition1, condition2)
     * Not	不满足条件	Filter.not(condition)
     * IsEqualTo	等于	new IsEqualTo("field", "value")
     * IsGreaterThan	大于	new IsGreaterThan("field", value)
     * IsLessThan	小于	new IsLessThan("field", value)
     * IsIn	在列表内	new IsIn("field", listOfValues)
     */
    @GetMapping("getFilterTextSegment")
    public void getFilterTextSegment(@RequestParam String collectionName){
        QdrantEmbeddingStore embeddingStore = QdrantEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6334)
                .collectionName(collectionName)
                .build();
        Embedding queryEmbedding = embeddingModel.embed("404 是哪里的问题？").content();


        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(metadataKey("author").isEqualTo("冷冷"))
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
//        System.out.println(searchResult.matches().get(0).score());  // 输出匹配分数
//        System.out.println(searchResult.matches().get(0).embedded().text());  // 输出匹配文本
        System.out.println(searchResult.matches().get(0).embedded().text());
    }

}
